#!/usr/bin/python

import sys,os,re
from pylab import *
from scipy.ndimage import filters,morphology,interpolation,measurements
from scipy import stats
import tables
from llpy.sutils import *

import argparse
parser = argparse.ArgumentParser(description="Different rescaling methods.")
parser.add_argument('input',default='data.h5',help="prototype file")
parser.add_argument('-d','--size',default=24,type=int,help='target size')
parser.add_argument('-o','--output',default='zoomed.h5',help="output file")
parser.add_argument('-D','--display',action="store_true",help="debugging display")
parser.add_argument('-N','--nsamples',type=int,default=2000000000,help="max # of samples")
parser.add_argument('-r','--pattern',default='.*',help='pattern selection (regex)')
args = parser.parse_args()

scale = None

with tables.openFile(args.input) as db, tables.openFile(args.output,"w") as odb:
    nsamples = min(args.nsamples,len(db.root.classes))
    ishape = list(db.root.patches.shape[1:])
    if scale is None:
        scale = args.size*1.0/ishape[0]
        oshape = list(interpolation.zoom(db.root.patches[0],(scale,scale),order=1).shape)
        print "zooming from",ishape,"to",oshape
    table_lcopy(db,odb)
    table_log(odb,str(sys.argv))
    rels = db.root.rel if "rel" in dir(db.root) else None
    odb.createEArray(odb.root,"patches",tables.Float32Atom(),shape=[0]+oshape,filters=tables.Filters(9))
    odb.createEArray(odb.root,"classes",tables.Int64Atom(),shape=(0,),filters=tables.Filters(9))
    if rels is not None:
        odb.createEArray(odb.root,"rel",tables.Float32Atom(),shape=(0,3),filters=tables.Filters(9))
    for i in range(nsamples):
        cls = udecode(db.root.classes[i])
        if not re.match(args.pattern,cls): continue
        p = db.root.patches[i]
        if i%1000==0: print i
        pdb = interpolation.zoom(p,(scale,scale),order=1)
        pdb = csnormalize(pdb)
        if args.display:
            ion(); clf(); gray(); imshow(pdb,interpolation='nearest'); ginput(1,0.001)
        odb.root.classes.append([db.root.classes[i]])
        odb.root.patches.append([pdb])
        if rels is not None:
            odb.root.rels.append([rels[i]])
